#!/usr/bin/env python
# malwarelu.py
#                  _                          _       
#  _ __ ___   __ _| |_      ____ _ _ __ ___  | |_   _ 
# | '_ ` _ \ / _` | \ \ /\ / / _` | '__/ _ \ | | | | |
# | | | | | | (_| | |\ V  V / (_| | | |  __/_| | |_| |
# |_| |_| |_|\__,_|_| \_/\_/ \__,_|_|  \___(_)_|\__,_|
#
# Script to make request on the malware.lu API
#
# It use multipartposthandler class available on last ubuntu
#   $ sudo apt-get install python-multipartposthandler
# if you don't have it you can grab it with easy_install
#   $ sudo easy_install MultipartPostHandler
#   $ easy_install --user MultipartPostHandler
# last solution take it online and put it on the same
# folder as the script
#
# Warning no SSL certificat verification
# due to urllib2, so don't use it on hostile
# connection, if you patch the code to support this feature
# contact us
#
# any subjection, tips, improvement are welcome

import re
import hashlib
import urllib, urllib2, simplejson
import argparse, re, sys
import MultipartPostHandler

apikey = ""
version = "0.6"

class MalwareLuAlreadyHere(Exception):
    def __init__(self, md5):
        self.md5 = md5
    def __str__(self):
        return "Hash %s already in database" % self.md5

class MalwareLu():

    def __init__(self, apikey):
        base = "https://www.malware.lu/api/"

        self.apikey = apikey

        self.url_download = base + "download"
        self.url_upload = base + "upload"
        self.url_check = base + "check" 
        self.url_stats = base + "stats"
        self.url_search = base + "search"
        self.url_submiturl = base + "submit"

        
    def filetohash(self, filename):
        file_content = open(filename, "rb").read()
        md5 = hashlib.md5(file_content).hexdigest()
        return md5

    def get_content_type(self, filename):
        return mimetypes.guess_type(filename)[0] or 'application/octet-stream'
     
    def download(self, md5):
        param = { 'hash': md5,
                'apikey': self.apikey }
        data = urllib.urlencode(param)
        req =  urllib2.Request(self.url_download, data)
        f = urllib2.urlopen(req)

        headers = f.info()
        if headers['content-type'] == 'application/octet-stream':
            filename = re.findall('filename=(\S+)', 
                headers['Content-Disposition'])[0]
            fp = open(filename, 'w')
            fp.write(f.read())
            fp.close()
            return True
        return simplejson.loads(f.read())

    def check(self, md5):
        param = { 'hash': md5,
                'apikey': self.apikey }
        data = urllib.urlencode(param)
        f = urllib2.urlopen(self.url_check, data)
        data = f.read()
        return simplejson.loads(data)
    
    def search(self, name):
        param = { 'name': name,
                'apikey': self.apikey }
        data = urllib.urlencode(param)
        f = urllib2.urlopen(self.url_search, data)
        data = f.read()
        return simplejson.loads(data)
    
    def stats(self):
        param = { 'apikey': self.apikey }
        data = urllib.urlencode(param)
        f = urllib2.urlopen(self.url_stats, data)
        data = f.read()
        return simplejson.loads(data)
        

    def submit_url(self, url):
        param = { 'url': url,
                'apikey': self.apikey }
        data = urllib.urlencode(param)
        f = urllib2.urlopen(self.url_submiturl, data)
        data = f.read()
        return simplejson.loads(data)


    def upload(self, filename):
        if self.check(self.filetohash(filename))['status'] == True:
            raise MalwareLuAlreadyHere(self.filetohash(filename))
        opener = urllib2.build_opener(MultipartPostHandler.MultipartPostHandler)
        params = {'file': open(filename, 'rb'), 'apikey': self.apikey}
        f = opener.open(self.url_upload, params)
        data = f.read()
        return simplejson.loads(data)



class bcolors:
    GREEN = '\033[92m'
    RED = '\033[91m'
    ENDC = '\033[0m'

    def disable(self):
        self.GREEN = ''
        self.RED = ''
        self.ENDC = ''

def show_sample_info(s):
    print "-"*72
    print "SHA256: %s" % s['sha256']
    print "MD5: %s" % s['md5']
    print ("Detection ratio: %s%s/%s%s") % \
        (bcolors.RED, s['positives'], s['total'], bcolors.ENDC)
    print ("Insert date: %s%s%s") % (bcolors.GREEN, s['date'],  
        bcolors.ENDC)
    print ("Scan date: %s%s%s") % (bcolors.GREEN, s['scan_date'],  
        bcolors.ENDC)
    print ("First seen: %s%s%s") % (bcolors.GREEN, s['first_seen'],
        bcolors.ENDC)
    print ("Last seen: %s%s%s") % (bcolors.GREEN, s['last_seen'],
        bcolors.ENDC)
    print "Size: %s" % s['file_size']

    av = ['nod32', 'avast', 'microsoft', 'clamav', 'kaspersky', 
            'bitdefender', 'sophos', 'avg', 'f_secure', 
            'mcafee_gw_edition', 'symantec', 'gdata']
    for a in av:
        if s[a] != None:
            print ("%s: %s%s%s") % (a, bcolors.RED, s[a], bcolors.ENDC)

    tools = ['peid', 'f_prot_unpacker' ]
    for t in tools:
        if s[t] != None:
            print ("%s: %s%s%s") % (t, bcolors.RED, s[t], bcolors.ENDC)


def main():
    parser = argparse.ArgumentParser(description = "Malware.lu API tools")
    parser.add_argument('--version', action='version', 
        version="%(prog)s version " + version)
    parser.add_argument('-k', '--apikey', action="store", default=apikey)

    group = parser.add_mutually_exclusive_group(required=True)

    group.add_argument('-c', '--check', nargs="+", dest="check", metavar="hash",
        help="Check a hash on malware.lu")
    group.add_argument('--stats', dest="action", 
        action='store_const', const="stats",
        help="Show your stats on malware.lu")
    group.add_argument('-s', '--search',
        help="Search by AV name on malware.lu")

    group.add_argument('--submit-url', nargs="+",
        help="Submit a malware by URL on malware.lu")

    group.add_argument('-u', '--upload', nargs="+", dest="upload", metavar="file",
        help="Upload a file to malware.lu")
    group.add_argument('-d', '--download', nargs="+", dest="download", metavar="hash",
        help="Download a hash from malware.lu")
    
    r = parser.parse_args()
    
    m = MalwareLu(r.apikey)
    try:
        # Upload files loop
        if r.upload != None:
            for arg in r.upload:
                print("Upload of %s..." % arg)
                try:
                    ret = m.upload(arg)
                    if ret['status'] == True :
                        print "Upload %s successfully" % arg
                    else:
                        print ret['error'] 
                except MalwareLuAlreadyHere as e:
                   print e


        # Download files loop
        if r.download != None:
            for arg in r.download:
                print "Download of %s..." % arg
                ret = m.download(arg)
                if ret == True :
                    print "Download %s successfully" % arg
                else:
                    print ret['error']
        
        # Check files loop
        if r.check != None:
            for arg in r.check:
                ret = m.check(arg)
                if ret['status'] == True:
                    print "%s is in database" % arg
                else:
                    print "%s not in database" % arg
        
        if r.submit_url != None:
            for arg in r.submit_url:
                ret = m.submit_url(arg)
                if ret['status'] == True:
                    print "URL %s correctly submit with md5 %s" % (arg, ret['md5'])
                else:
                    print ret['error']

        if r.search != None:
            ret = m.search(r.search)
            for s in ret['samples']:
                show_sample_info(s)

        # stats action
        if r.action == "stats":
            ret = m.stats()
            print "Downloads stats %s" % ret['stats']

    except urllib2.HTTPError, e:
        if e.code == 401:
            print "Wrong api key"
        else:
            print e

if __name__ == '__main__':
	main()
